# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.
# Adapted from https://github.com/MARIO-Math-Reasoning/Super_MARIO
"""
rStar-Math 核心推理求解器 (Solver)

这个模块是 rStar-Math 系统的核心控制器，负责协调整个推理过程：
1. 管理 LLM 和奖励模型的初始化
2. 控制 MCTS/束搜索的执行流程
3. 协调生成、评估、选择的三个阶段
4. 管理多轮 rollout 和中间结果保存

主要组件：
- Solver: 核心求解器类，协调所有组件
- generate_preprocess/postprocess: 处理 LLM 生成阶段
- value_preprocess/postprocess: 处理价值评估阶段
- save_intermediate_*: 保存中间结果用于分析

典型工作流程：
1. 初始化 → 2. 多轮rollout → 3. 每轮多步推理 → 4. 生成候选 → 5. 评估价值 → 6. 选择下一步
"""

from __future__ import annotations  # 启用延迟注解评估，支持前向引用类型注解
import os  # 操作系统接口，用于文件路径操作和环境变量设置
import json  # JSON数据格式处理，用于序列化和反序列化数据
import os.path as osp  # 路径操作工具，提供跨平台的文件路径处理功能
from tqdm import tqdm  # 进度条显示库，用于显示长时间运行任务的进度
from termcolor import colored  # 终端彩色输出库，用于在控制台显示彩色文本
from functools import partial  # 函数工具库，用于创建偏函数（固定部分参数的函数）
from vllm import LLM, SamplingParams  # vLLM推理引擎，用于高效的大语言模型推理和采样参数配置
from vllm.outputs import RequestOutput  # vLLM输出类型，表示推理请求的输出结果
from pebble import ProcessPool  # 进程池库，用于并行处理和多进程任务管理
from omegaconf import DictConfig, OmegaConf  # 配置管理库，用于处理结构化配置和YAML文件
from typing import Optional, Any, Dict, List, Callable, Type, Tuple  # 类型注解工具，用于静态类型检查和代码文档
from pydantic import BaseModel, ConfigDict, field_validator  # 数据验证库，用于数据模型定义和字段验证
from .agents.tree import BaseTree  # 基础树结构，用于搜索算法的树形数据结构
from .agents.mcts import MCTS  # 蒙特卡洛树搜索算法实现
from .llms.llms import llm_generate, rm_generate  # LLM生成函数：语言模型生成和奖励模型生成
from .llms.llm_engine import llm_engine, rm_engine  # LLM引擎：语言模型引擎和奖励模型引擎
from .constants import TIMEOUT_SECONDS, ERROR_COLOR  # 常量定义：超时时间和错误颜色配置


class Solver(BaseModel):
    """
    rStar-Math 核心推理求解器
    
    作为整个推理系统的中央控制器，协调 LLM 生成、搜索算法、奖励模型评估等各个组件。
    支持 MCTS 和束搜索两种推理模式，通过多轮 rollout 实现深度推理。
    
    主要职责：
    1. 初始化和管理 LLM 推理引擎
    2. 可选地初始化 PPM 奖励模型
    3. 控制多轮推理的执行流程
    4. 协调生成-评估-选择的完整循环
    5. 保存中间结果用于分析和调试
    
    属性说明：
        config: 系统配置对象，包含所有推理参数
        llm: LLM 生成函数的部分应用
        reward_model: PPM 奖励模型函数的部分应用（可选）
        need_value_func: 是否启用价值函数评估
        max_agent_steps: 最大代理步骤数（rollout 次数）
    """
    # Pydantic 配置：允许任意类型的属性
    model_config = ConfigDict(arbitrary_types_allowed=True)
    
    # 核心配置和模型组件
    config: Any  # 系统配置对象
    stop: List[str] = None  # 生成停止标记列表
    llm: Optional[Callable[[...], List[str]]] = None  # LLM 生成函数
    llm_engine: Optional[LLM] = None  # vLLM 推理引擎实例
    generate_sampling_params: Optional[SamplingParams] = None  # 生成采样参数
    need_value_func: bool = False  # 是否需要价值函数评估
    max_agent_steps: int = 1  # 最大代理步骤数（rollout次数）
    reward_model: Optional[Any] = None  # PPM 奖励模型函数

    def __init__(self, **kwargs) -> None:
        """
        初始化求解器，设置 LLM 和可选的奖励模型
        
        根据配置初始化必要的组件：
        1. 解析停止标记配置
        2. 根据 need_value_func 标志决定是否创建奖励模型
        3. 创建 LLM 推理引擎
        4. 根据推理模式设置最大步骤数
        
        参数:
            **kwargs: 传递给父类的初始化参数，主要是 config
        """
        super().__init__(**kwargs)
        
        # === Caption 模式配置验证 ===
        # 在 Caption 模式下，确保 n_generate_sample 与 reason_samples + caption_samples 一致
        if self.config.mode == "mcts_caption":
            expected_n_generate_sample = self.config.reason_samples + self.config.caption_samples
            if self.config.n_generate_sample != expected_n_generate_sample:
                print(f"🔧 Caption 模式：n_generate_sample ({self.config.n_generate_sample}) 与 reason_samples + caption_samples ({expected_n_generate_sample}) 不一致，已自动调整")
                self.config.n_generate_sample = expected_n_generate_sample
        
        # 解析停止标记配置
        if self.config.stop:
            self.stop = OmegaConf.to_object(self.config.stop)

        # 设置是否需要价值函数评估
        self.need_value_func = self.config.need_value_func
        if self.need_value_func:
            # 如果启用价值函数，则创建 PPM 奖励模型
            self.reward_model = self.create_rm()
            
        # 创建 LLM 生成函数
        self.llm = self.create_llm()
        
        # 根据推理模式设置最大代理步骤数
        if self.config.mode == "sbs":  # Step-by-step 模式
            self.max_agent_steps = 1
        elif self.config.mode in ["mcts", "mcts_wo_code", "mcts_caption"]:  # MCTS 模式
            self.max_agent_steps = self.config.iterations  # 使用配置的迭代次数
            self.config.step_beam_width = 1  # MCTS 模式下束宽度固定为1
        else:
            self.max_agent_steps = self.config.iterations


    @field_validator("config")
    def validate_config(cls, cfg: Any):
        """
        验证配置对象的类型
        
        确保传入的配置对象是 DictConfig 类型的子类，这是 OmegaConf 的要求。
        
        参数:
            cfg: 待验证的配置对象
            
        返回:
            验证通过的配置对象
            
        异常:
            TypeError: 如果配置对象类型不正确
        """
        if issubclass(type(cfg), DictConfig):
            return cfg
        raise TypeError("Wrong type for `config`, must be subclass of BaseConfig")


    def create_rm(self):
        """
        创建 PPM（过程偏好模型）奖励模型
        
        初始化奖励模型引擎，包括模型、价值头和分词器，然后创建一个部分应用的
        生成函数以便后续调用。奖励模型用于评估推理步骤的质量。
        
        返回:
            functools.partial: 配置好的奖励模型生成函数
            
        依赖:
            - config.reward_model_dir: 奖励模型路径
            - config.max_model_len: 模型最大序列长度
        """
        # 使用 rm_engine 创建奖励模型、价值头和分词器
        rm, v_head, tokenizer = rm_engine(self.config)
        
        # 创建部分应用的奖励模型生成函数
        return partial(
            rm_generate,  # 奖励模型生成函数
            model=rm,     # 奖励模型实例
            v_head=v_head,  # 价值头网络
            tokenizer=tokenizer,  # 分词器
            max_model_len=self.config.max_model_len,  # 最大序列长度
        )


    def create_llm(self):
        """
        创建 LLM（大语言模型）推理引擎
        
        初始化 vLLM 推理引擎和采样参数，然后创建一个部分应用的生成函数。
        这个函数将用于生成推理步骤的候选解。
        
        返回:
            functools.partial: 配置好的 LLM 生成函数
            
        副作用:
            设置 self.llm_engine 和 self.generate_sampling_params
        """
        # 使用 llm_engine 创建推理引擎和采样参数
        engine, sampling_params = llm_engine(self.config)
        
        # 保存引擎实例和采样参数供后续使用
        self.llm_engine = engine
        self.generate_sampling_params = sampling_params
        
        # 为OpenAI引擎添加配置引用，用于多模态数据处理
        if self.config.backend_type == "openai":
            engine._config = self.config
        
        # 创建部分应用的 LLM 生成函数
        return partial(
            llm_generate,  # LLM 生成函数
            engine=self.llm_engine,  # vLLM 推理引擎
        )

        
    @staticmethod
    def processor(agent, output) -> BaseTree:
        """
        静态方法：处理单个代理的 LLM 输出
        
        调用代理的 generate_next_step 方法来处理 LLM 生成的输出，
        包括解析文本、执行代码、创建子节点等操作。
        
        参数:
            agent: 搜索代理实例（MCTS 或 BS）
            output: LLM 生成的输出结果
            
        返回:
            BaseTree: 处理后的代理对象
            
        注意:
            这是一个静态方法，用于多进程池处理
        """
        agent.generate_next_step(output)
        return agent


    @staticmethod
    def selector(agent, output) -> BaseTree:
        """
        静态方法：处理单个代理的价值评估和选择
        
        调用代理的 select_next_step 方法来处理奖励模型的输出，
        更新节点价值并选择下一个要探索的节点。
        
        参数:
            agent: 搜索代理实例（MCTS 或 BS）
            output: 奖励模型的评估输出
            
        返回:
            BaseTree: 处理后的代理对象
            
        注意:
            这是一个静态方法，用于价值评估阶段
        """
        agent.select_next_step(output)
        return agent


    def generate_preprocess(self, agents):
        """
        生成阶段的预处理：收集需要 LLM 生成的提示词和多模态数据
        
        遍历所有代理，收集需要生成下一步的代理的提示词，并按类型分组。
        这个方法确保只有需要生成的代理才会被处理，提高效率。
        
        特别优化：支持MCTSCaption的批量推理优化，根据prompt元数据动态调整采样参数。
        
        参数:
            agents: 所有搜索代理的列表
            
        返回:
            tuple: 包含以下元素的元组
                - prompts: 所有需要生成的提示词列表
                - prompts_span: 提示词在各代理中的分割位置
                - valid_agents: 需要生成的有效代理列表
                - invalid_agents: 不需要生成的无效代理列表
                - expanded_agents: 已经扩展过的代理列表
                - rewards: 已有的奖励值列表
                - multimodal_data: 多模态数据列表
                
        工作流程:
        1. 遍历所有代理，检查是否需要生成下一步
        2. 对于需要生成的代理，区分已扩展和未扩展
        3. 收集未扩展代理的提示词和奖励信息
        4. 收集多模态数据（如果启用）
        5. 记录提示词在各代理中的边界位置
        6. 特殊处理MCTSCaption：动态调整采样参数以支持批量推理优化
        """
        prompts = []
        rewards = []  # 所有奖励值的列表
        prompts_span = [0]  # 提示词在各代理中的边界索引
        valid_agents = []    # 需要生成且未扩展的代理
        invalid_agents = []  # 不需要生成的代理
        expanded_agents = [] # 已经扩展过的代理
        multimodal_data = [] # 多模态数据列表

        for agent in agents:
            if agent.should_generate_next():  # 检查是否需要生成下一步
                if agent.has_expanded():  # 已经扩展过的代理
                    expanded_agents.append(agent)
                else:  # 未扩展的代理，需要收集提示词
                    agent_prompts = agent.create_prompt()  # 创建提示词
                    rewards.extend(agent.get_rewards())    # 收集奖励信息
                    prompts.extend(agent_prompts)          # 收集提示词
                    
                    # 收集多模态数据
                    if self.config.multimodal_enabled and hasattr(agent, 'root') and agent.root.multimodal_data:
                        # 为每个提示词复制多模态数据
                        for _ in range(len(agent_prompts)):
                            multimodal_data.append(agent.root.multimodal_data)
                    else:
                        # 为每个提示词添加空的多模态数据
                        for _ in range(len(agent_prompts)):
                            multimodal_data.append(None)
                    
                    # 多模态数据处理完成，无需特殊的MCTSCaption处理
                    
                    # 记录该代理提示词的结束位置
                    prompts_span.append(prompts_span[-1] + len(agent_prompts))
                    valid_agents.append(agent)
            else:  # 不需要生成的代理
                invalid_agents.append(agent)
        
        return prompts, prompts_span, valid_agents, invalid_agents, expanded_agents, rewards, multimodal_data


    def generate_postprocess(
        self, 
        outputs: List[List[RequestOutput]], 
        valid_agents: List[BaseTree],
    ) -> List[BaseTree]:
        """
        生成阶段的后处理：并行处理 LLM 输出结果
        
        使用多进程池并行处理所有有效代理的 LLM 输出，包括：
        1. 解析 LLM 生成的文本内容
        2. 执行其中的 Python 代码
        3. 创建新的子节点
        4. 处理异常和超时情况
        
        参数:
            outputs: LLM 生成的输出结果列表，每个代理对应一个子列表
            valid_agents: 需要处理的有效代理列表
            
        返回:
            List[BaseTree]: 处理后的代理列表，失败的代理保持原状
            
        注意:
            - 使用固定 12 个工作进程的进程池
            - 设置超时保护避免无限等待
            - 异常处理确保程序稳定性
        """
        post_agents = []
        # 创建进程池，使用配置的并行工作进程数
        max_workers = min(self.config.step_parallel_workers, len(valid_agents), os.cpu_count())
        with ProcessPool(max_workers=max_workers) as pool:
            # 提交任务到进程池，设置超时限制
            future = pool.map(self.__class__.processor, valid_agents, outputs, timeout=TIMEOUT_SECONDS)
            iterator = future.result()
        
        # 创建进度条显示处理进度
        progress_bar = tqdm(total=len(valid_agents), desc="generate_postprocess")  
        while True:
            try:
                result = next(iterator)  # 获取下一个处理结果
                post_agents.append(result)
            except StopIteration:
                break  # 所有任务处理完成
            except Exception as error:
                # 处理异常，记录错误信息并继续
                print(colored(f"{error}\n", ERROR_COLOR))
                post_agents.append(None)  # 标记失败的任务
            progress_bar.update(1)  # 更新进度条
        progress_bar.close()  # 关闭进度条
            
        # 更新代理列表：成功的使用新结果，失败的保持原状
        updated_agents = [
            post_agent if post_agent is not None else valid_agent
            for post_agent, valid_agent in zip(post_agents, valid_agents)
        ]
        return updated_agents
    

    def value_preprocess(self, agents: List[BaseTree]) -> Tuple[List[str], List[int]]:
        """
        价值评估阶段的预处理：收集需要评估的提示词
        
        为所有代理创建价值评估提示词，用于 PPM 奖励模型评估推理步骤的质量。
        与 generate_preprocess 不同，这里对所有代理都进行处理。
        
        参数:
            agents: 所有需要评估的搜索代理列表
            
        返回:
            tuple: 包含以下元素的元组
                - prompts: 所有价值评估提示词的列表
                - prompts_span: 提示词在各代理中的分割位置
                
        工作流程:
        1. 遍历所有代理
        2. 为每个代理创建价值评估提示词
        3. 记录每个代理提示词的边界位置
        """
        prompts = []  # 所有价值评估提示词
        prompts_span = [0]  # 提示词在各代理中的边界索引
        
        for agent in agents:
            # 为每个代理创建价值评估提示词
            agent_prompts = agent.create_prompt(is_value_only=True)
            prompts.extend(agent_prompts)
            # 记录该代理提示词的结束位置
            prompts_span.append(prompts_span[-1] + len(agent_prompts))
        return prompts, prompts_span
    
    
    def value_postprocess(
        self, 
        outputs, 
        valid_agents,
    ) -> List[BaseTree]:
        """
        价值评估阶段的后处理：处理奖励模型输出并选择下一步
        
        使用奖励模型的输出来更新节点价值，并选择下一个要探索的节点。
        这个阶段是 MCTS 算法中的“选择”和“回溯”阶段。
        
        参数:
            outputs: 奖励模型的评估输出结果列表
            valid_agents: 需要处理的有效代理列表
            
        返回:
            List[BaseTree]: 处理后的代理列表
            
        工作流程:
        1. 遍历所有有效代理
        2. 对于非空的代理，调用 selector 处理奖励输出
        3. 更新节点价值并选择下一个探索目标
        """
        for agent, output in zip(valid_agents, outputs):
            if agent is not None:  # 只处理非空的代理
                # 调用选择器处理奖励模型输出，更新价值并选择下一步
                self.selector(agent, output)
        return valid_agents


    def save_intermediate_metric(self, path: str, agents: List[MCTS], rollout) -> None:
        """
        保存中间指标数据用于分析和调试
        
        计算并保存每个 rollout 的性能指标，包括 pass@1 和 pass@n 的准确率。
        这些指标可以帮助分析 MCTS 算法的收敛过程和效果。
        
        参数:
            path: 保存文件的路径
            agents: MCTS 代理列表
            rollout: 当前 rollout 的索引
            
        注意:
            - 只在非采样模式下执行（is_sampling=False）
            - pass@1: 最高价值的解是否正确
            - pass@n: 是否存在任意一个正确的解
        """
        if self.config.is_sampling: return  # 采样模式下不保存
        
        # 提取所有代理的中间指标
        states = [s.intermediate_metric for s in agents]
        statics = []  # 统计结果列表
        
        # 对每个 rollout 计算指标
        for i in range(rollout + 1):
            pass1, passn = 0, 0  # pass@1 和 pass@n 的计数
            
            # 遍历每个代理的状态
            for idx, state in enumerate(states):
                max_value = -100  # 最大价值
                max_value_result = False  # 最大价值对应的结果
                pass1_ans = False  # 是否有正确答案
                
                # 检查该 rollout 索引之前的所有结果
                for idx, rollout_index in enumerate(state["rollout_indexs"]):
                    if rollout_index <= i:  # 只考虑当前 rollout 之前的结果
                        # 更新最大价值和对应结果
                        if state["value_estimate"][idx] > max_value:
                            max_value = state["value_estimate"][idx]
                            max_value_result = state["judgements"][idx]
                        # 检查是否有正确答案
                        if state["judgements"][idx]:
                            pass1_ans = True
                            
                # 统计 pass@1 和 pass@n
                if max_value_result:  # 最高价值的解正确
                    pass1 += 1
                if pass1_ans:  # 存在正确解
                    passn += 1
                    
            # 保存该 rollout 的统计结果
            statics.append({
                "rollout": i,
                "pass1": pass1,    # pass@1 正确数
                "passn": passn,    # pass@n 正确数
                "len": len(states), # 总问题数
            })
            
        # 将统计结果和详细状态保存到文件
        with open(path, "w", encoding='utf-8') as f:
            json.dump([statics, states], f, ensure_ascii=False, indent=4)

    
    def save_intermediate_rollouts(self, saved_jsonl_file, cur_data, agents, rollout_idx):
        """
        保存中间 rollout 结果用于分析和调试
        
        如果启用了中间结果保存，将当前 rollout 的结果保存到单独的文件中。
        这对于分析 MCTS 的收敛过程和调试算法参数非常有用。
        
        参数:
            saved_jsonl_file: 主输出文件的路径
            cur_data: 当前批次的问题数据
            agents: 所有搜索代理列表
            rollout_idx: 当前 rollout 的索引
            
        功能:
        1. 保存中间指标数据到 rollout 目录
        2. 保存当前 rollout 的详细结果
        3. 创建必要的目录结构
        """
        # 检查是否需要保存中间结果
        if self.config.save_intermediate_rollouts and saved_jsonl_file and (self.config.mode == "mcts" or self.config.mode == "mcts_wo_code" or self.config.mode == "mcts_caption"):
            # 解析输出文件路径，使用新的目录结构
            saved_json_dir = os.path.dirname(os.path.dirname(saved_jsonl_file))  # 回到run_dir
            rollout_dir = os.path.join(saved_json_dir, "rollouts")
            saved_jsonl_file_name = os.path.basename(saved_jsonl_file)
            
            # 确保rollout目录存在
            if not os.path.exists(rollout_dir):
                os.makedirs(rollout_dir)
                
            # 保存中间指标数据
            self.save_intermediate_metric(
                os.path.join(rollout_dir, f"intermediate_metric_{saved_jsonl_file_name}"), 
                agents, 
                rollout_idx
            )
            
            # 获取当前 rollout 的输出结果
            outs = self.output(agents)
            
            # 保存当前 rollout 的详细结果
            rollout_file = os.path.join(rollout_dir, f"rollout{rollout_idx:02d}_{saved_jsonl_file_name}")
            with open(rollout_file, "a+", encoding='utf-8') as writer:
                for d in cur_data:
                    question = d["question"]
                    d["rstar"] = outs[question]  # 添加 rStar 推理结果
                    writer.write(json.dumps(d, ensure_ascii=False) + '\n')
                    writer.flush()
    
    def output(self, agents: List[BaseTree]):
        """
        生成最终输出结果
        
        从所有代理中提取最终的推理结果，并组织成字典格式。
        每个问题对应一个完整的推理过程和最终答案。
        
        参数:
            agents: 所有搜索代理的列表
            
        返回:
            dict: 以问题为键，推理结果为值的字典
            
        注意:
            - 调用每个代理的 return_states() 方法
            - 返回的结果包含完整的推理路径和最终答案
        """
        jsonlines = {}
        # 遍历所有代理，提取其推理结果
        for i, agent in enumerate(agents):         
            jsonlines[agent.question] = agent.return_states()
        
        return jsonlines
    
    def solve(self, agents: List[BaseTree], saved_jsonl_file: str, cur_data: List[Dict[str, Any]]):
        """
        核心求解方法：执行完整的推理过程
        
        这是 Solver 的核心方法，实现了完整的 rStar-Math 推理流程：
        1. 多轮 rollout 执行（MCTS 中的多次模拟）
        2. 每轮内的多步推理（探索树的扩展）
        3. 生成-评估-选择的循环流程
        4. 中间结果的保存和跟踪
        
        参数:
            agents: 所有问题对应的搜索代理列表
            saved_jsonl_file: 输出文件路径，用于中间结果保存
            cur_data: 当前批次的问题数据
            
        返回:
            dict: 最终的推理结果，以问题为键
            
        核心流程:
        1. 外层循环: rollout 迭代（对于 MCTS，这是多次独立的模拟）
        2. 内层循环: 每轮内的推理步骤（探索树的逐步扩展）
        3. 每步包含: 预处理 → LLM生成 → 后处理 → 价值评估 → 选择下一步
        """
        
        print(f"\n🚀 =============== 开始推理求解 ===============")
        print(f"📊 总共 {len(agents)} 个问题，每个问题将执行 {self.max_agent_steps} 轮搜索")
        print(f"🔧 算法模式：{self.config.mode}")
        print(f"🎯 每轮最大搜索深度：{self.config.max_depth}")
        if hasattr(self.config, 'reason_samples') and hasattr(self.config, 'caption_samples'):
            print(f"🎭 动作配置：Reason={self.config.reason_samples}, Caption={self.config.caption_samples}")
        
        # === 外层循环：多轮 rollout 执行 ===
        for rollout in tqdm(range(self.max_agent_steps), desc="🔄 搜索轮次进度", 
                           bar_format='{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]'):
            # 初始化每轮 rollout 的起始搜索点，每轮都从根节点开始
            for agent in agents:
                agent.select_next_step(from_root=True)  # 重置到根节点
                agent.rollout_idx = rollout  # 设置当前 rollout 索引

            # === 内层循环：每轮内的多步推理 ===
            print(f"\n🎯 =============== 第 {rollout + 1}/{self.max_agent_steps} 轮搜索 ===============")
            print(f"📝 说明：每个问题都有独立的搜索树，本轮将继续扩展所有问题的搜索树")
            
            for step in range(self.config.max_depth):
                print(f"\n🔍 --- 搜索步骤 {step + 1}/{self.config.max_depth} ---")
                
                # === 阶段1：生成阶段预处理 ===
                # 收集需要 LLM 生成的提示词和代理信息
                prompts, prompts_span, valid_agents, invalid_agents, expanded_agents, valid_rewards, multimodal_data = self.generate_preprocess(agents)
                
                # 如果没有需要处理的代理，结束当前 rollout
                if len(valid_agents + expanded_agents) < 1:
                    print("⏹️  所有问题的搜索树都已完成扩展，提前结束本轮搜索")
                    break
                
                # 解释agent的含义
                total_active = len(valid_agents) + len(expanded_agents)
                print(f"📊 搜索状态统计：")
                print(f"   • 需要LLM生成的问题：{len(valid_agents)} 个")
                print(f"   • 已完成生成待评估的问题：{len(expanded_agents)} 个") 
                print(f"   • 已终止搜索的问题：{len(invalid_agents)} 个")
                print(f"   • 总活跃问题数：{total_active} 个")
                
                if len(valid_agents) > 0:
                    print(f"🤖 准备调用LLM生成 {len(prompts)} 个推理候选")
                    
                    # === 调试输出：记录发送的prompts ===
                    if hasattr(self.config, 'debug_prompts') and self.config.debug_prompts:
                        debug_dir = os.path.join(os.path.dirname(saved_jsonl_file), "..", "logs")
                        os.makedirs(debug_dir, exist_ok=True)
                        
                        debug_file = os.path.join(debug_dir, f"prompts_rollout{rollout}_step{step}.txt")
                        with open(debug_file, "w", encoding='utf-8') as f:
                            f.write(f"=== Rollout {rollout+1}, Step {step+1} ===\n")
                            f.write(f"Total prompts: {len(prompts)}\n")
                            f.write(f"Valid agents: {len(valid_agents)}\n\n")
                            
                            for i, prompt in enumerate(prompts):
                                f.write(f"--- Prompt {i+1} ---\n")
                                f.write(prompt)
                                f.write("\n\n")
                        
                        print(f"📝 Prompt调试信息已保存到: {debug_file}")
                    
                    # === 阶段2： LLM 生成阶段 ===
                    # 使用 LLM 生成候选推理步骤
                    outputs = self.llm(prompts, self.generate_sampling_params, multimodal_data=multimodal_data)
                    
                    # === 调试输出：记录LLM的输出结果 ===
                    if hasattr(self.config, 'debug_prompts') and self.config.debug_prompts:
                        debug_file = os.path.join(debug_dir, f"outputs_rollout{rollout}_step{step}.txt")
                        with open(debug_file, "w", encoding='utf-8') as f:
                            f.write(f"=== LLM Outputs - Rollout {rollout+1}, Step {step+1} ===\n")
                            f.write(f"Total outputs: {len(outputs)}\n\n")
                            
                            for i, output in enumerate(outputs):
                                f.write(f"--- Output {i+1} ---\n")
                                f.write(f"Number of completions: {len(output.outputs)}\n")
                                for j, completion in enumerate(output.outputs):
                                    f.write(f"Completion {j+1}:\n")
                                    f.write(f"Text: {completion.text}\n")
                                    f.write(f"Stop reason: {completion.stop_reason}\n")
                                    f.write("---\n")
                                f.write("\n")
                        
                        print(f"📝 输出调试信息已保存到: {debug_file}")
                    
                    # 将奖励信息附加到输出上，避免重复评估
                    for output, reward in zip(outputs, valid_rewards):
                        output.value_estimate = reward
                        
                    # 按代理分组重新组织输出结果
                    reconstructed_outputs = [outputs[bos_idx : eos_idx] for bos_idx, eos_idx in zip(prompts_span, prompts_span[1:])]
                    
                    print(f"✅ LLM生成完成，共获得 {len(outputs)} 个推理候选")
                    
                    # === 阶段3：生成阶段后处理 ===
                    # 并行处理 LLM 输出，执行 Python 代码并创建子节点
                    print(f"⚙️  正在并行处理LLM输出（创建搜索树节点）...")
                    valid_agents = self.generate_postprocess(reconstructed_outputs, valid_agents)

                # === 阶段4：价值评估阶段预处理 ===
                # 为所有代理创建价值评估提示词
                all_agents_for_eval = valid_agents + expanded_agents
                if len(all_agents_for_eval) > 0:
                    prompts, prompts_span = self.value_preprocess(all_agents_for_eval)
                    
                    if self.need_value_func:  # 如果启用了奖励模型
                        print("🏆 使用PPM奖励模型评估节点价值...")
                        # 使用 PPM 奖励模型评估步骤价值
                        outputs = self.reward_model(prompts=prompts)
                        # 按代理分组重新组织评估结果
                        reconstructed_outputs = [outputs[bos_idx : eos_idx] for bos_idx, eos_idx in zip(prompts_span, prompts_span[1:])]
                    else:
                        print("⚡ 跳过价值评估（未启用奖励模型）")
                        # 不使用奖励模型时，设置为 None
                        reconstructed_outputs = [None] * (len(prompts_span) - 1)
                    
                    # === 阶段5：价值评估和选择阶段 ===
                    # 处理奖励模型输出，更新节点价值并选择下一步
                    print("🎯 更新节点价值并选择下一步搜索节点...")
                    valid_agents = self.value_postprocess(reconstructed_outputs[:len(valid_agents)], valid_agents)
                    # 对已扩展的代理只进行选择步骤，不进行价值评估
                    expanded_agents = self.value_postprocess(reconstructed_outputs[len(valid_agents):], expanded_agents)
                
                # === 阶段6：合并所有代理 ===
                # 将各类代理合并回原始列表
                agents = valid_agents + invalid_agents + expanded_agents
                
                print(f"✅ 步骤 {step + 1} 完成，继续处理 {len(agents)} 个问题的搜索树")

            # === 保存当前 rollout 的中间结果 ===
            if self.config.save_intermediate_rollouts:
                print(f"💾 保存第 {rollout + 1} 轮搜索的中间结果...")
                self.save_intermediate_rollouts(saved_jsonl_file, cur_data, agents, rollout)
            print(f"✅ 第 {rollout + 1} 轮搜索完成")
            
        print(f"\n🎉 所有 {self.max_agent_steps} 轮搜索执行完成！")
        print("📋 正在整理最终结果...")
        
        # === 返回最终结果 ===
        return self.output(agents)
    
    
